#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
"""
@File    :   tar.py
@Time    :   2024-05-22 09:45
@Author  :   坐公交也用券
@Version :   1.0
@Contact :   faith01238@hotmail.com
@Homepage : https://liumou.site
@Desc    :   OpenSSH安装器 - 支持多种压缩格式(zip/tar.gz)的OpenSSH安装、升级和管理
              功能包括：
              - 自动下载或使用本地OpenSSH压缩包
              - 支持多种压缩格式验证和解压
              - 完整的系统服务管理
              - 备份恢复机制
              - 自动创建必要符号链接
"""
import logging
import os
import subprocess
import sys
import time
import tarfile
import zipfile
from argparse import ArgumentParser
from datetime import datetime
from subprocess import getoutput, getstatusoutput
from sys import exit

# 创建一个自定义的日志记录器
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

# 创建一个自定义的日志处理器，设置其输出格式
formatter = logging.Formatter('%(asctime)s [%(levelname)s] %(filename)s:%(lineno)d | %(funcName)s | %(message)s')
handler = logging.StreamHandler(sys.stdout)
handler.setFormatter(formatter)
logger.addHandler(handler)


class OpenSSHInstaller:
	def __init__(self, urls, file_name):
		"""
		:param file_name: 文件名
		:param urls: 下载地址
		"""
		self.url = urls
		self.filename = file_name
		self.local_file = False
		# 记录旧版SSH版本
		self.old_version = None
		# 设置安装路径,默认: /usr/local/openssh
		self.ssh_install_dst = "/usr/local/openssh"
		self.ssh_dst_bin = os.path.join(self.ssh_install_dst, "bin")
		self.ssh_dst_sbin = os.path.join(self.ssh_install_dst, "sbin")
		self.ssh_dst_etc = os.path.join(self.ssh_install_dst, "etc")
		# 备份目录
		self.backup_dir = None
		# 设置服务名称
		self.service_old_name = 'sshd'
		self.service_new_name = 'sshd'
		# 定义必需文件列表
		self.required_files = [
			"bin/ssh", "bin/scp", "bin/sftp",
			"etc/sshd_config", "etc/ssh_config",
			"sbin/sshd"
		]

	def check_os(self):
		"""
		检查操作系统
		:return:
		"""
		res = subprocess.getstatusoutput("cat /etc/os-release| grep ^NAME")
		if res[0] != 0:
			logger.error("无法获取操作系统信息")
			exit(1)
		tmp = res[1].split("=")[1].replace('"', "").lower()
		release_ = tmp
		tmps = tmp.split(" ")
		if len(tmps) > 1:
			release_ = tmps[0]
		# 如果文件名称存在 kylin和arm64则进行系统检查，否则直接返回
		if release_ not in str(self.filename).lower():
			logger.error(f"当前操作系统: {release_} 与 文件名称不匹配,请检查文件名称是否正确")
			exit(1)
		logger.info(f"当前操作系统: {release_} 关键词匹配成功: [ {self.filename} ]")

	def link_ssh_conf(self):
		"""
		链接sshd_config文件
		:return:
		"""
		dst = "/etc/ssh"
		# 获取当前时间,格式如下： 202405201230
		formatted_datetime = datetime.now().strftime('%Y%m%d%H%M%S')
		# 拼接文件名
		conf_new = f"{dst}_{formatted_datetime}"
		# 开始备份
		if os.path.exists(dst):
			c = f"mv {dst} {conf_new}"
			if os.system(c) == 0:
				logger.info(f"备份成功: [ {c} ]")
			else:
				logger.error(f"备份失败: [ {c} ]")
				exit(1)

		os.symlink(self.ssh_dst_etc, dst)

	def delete_init_d(self):
		"""
		删除init.d文件
		:return:
		"""
		fl = ["/etc/init.d/sshd", "/etc/rc.d/init.d/sshd"]
		for file_ in fl:
			if os.path.exists(file_):
				os.remove(file_)

	def service(self, p: str):
		"""
		创建service
		:param p: 可执行文件绝对路径
		:return:
		"""
		file_ = f'/usr/lib/systemd/system/{self.service_new_name}.service'
		# conf = os.path.join(self.ssh_dst_etc, "sshd_config")
		if os.path.exists("/etc/init.d/sshd"):
			os.system('rm -f /etc/init.d/sshd')
		os.system('rm -f /etc/systemd/system/ssh.service')
		os.system('rm -f /etc/systemd/system/sshd.service')
		txt = f"""[Unit]
Description=This is a manually generated SSH service file
After=network.target
StartLimitIntervalSec=5

[Service]
Type=simple
User=root
ExecStart={p}
StandardOutput=journal
KillMode=control-group
Restart=on-failure
RestartPreventExitStatus=255
StandardError=inherit
SyslogIdentifier=sshd

[Install]
WantedBy=multi-user.target
Alias={self.service_new_name}.service\n"""
		try:
			with open(file_, 'w', encoding="utf-8") as fp:
				fp.write(txt)
		except Exception as e:
			logger.error(e)
			return
		os.system(f"systemctl daemon-reload")
		os.system(f"systemctl enable {self.service_new_name}.service")
		os.system(f"systemctl start {self.service_new_name}.service")
		os.system(f"systemctl status {self.service_new_name}.service")

	def get_version(self):
		"""
		获取旧版SSH版本
		:return:
		"""
		self.old_version = getoutput("ssh -V")

	def init_arg(self):
		"""
		初始化参数
		:return:
		"""
		# 获取文件名
		self.filename = self.url.split("/")[-1]
		# 设置bin
		self.ssh_dst_bin = os.path.join(self.ssh_install_dst, "bin")
		# 设置sbin
		self.ssh_dst_sbin = os.path.join(self.ssh_install_dst, "sbin")

	def download(self):
		"""
		下载文件
		:return:
		"""
		logger.info("正在下载文件: {}".format(self.filename))
		# 通过wget下载文件
		c = f"wget {self.url} -O {self.filename}"
		# 执行命令并判断是否下载成功，如果失败则终止程序
		if os.system(c) != 0:
			logger.error("下载失败，请检查网络")
			print(c)
			exit(2)
		logger.info(f"下载成功: [ {self.filename} ]")

	def extract_zip_archive(self):
		"""
		解压zip格式的归档文件
		:return:
		"""
		# 检查压缩包结构
		has_openssh = False
		required_files = self.required_files
		try:
			with zipfile.ZipFile(self.filename) as zf:
				# 检查是否包含openssh目录
				for name in zf.namelist():
					if name.startswith("openssh/") or name.startswith("./openssh/"):
						has_openssh = True
						break
				
				# 检查是否包含必要的文件
				missing_files = []
				for req_file in required_files:
					found = False
					for name in zf.namelist():
						if (has_openssh and (name == f"openssh/{req_file}" or name == f"./openssh/{req_file}")) or \
						   (not has_openssh and (name == req_file or name == f"./{req_file}")):
							found = True
							break
					if not found:
						missing_files.append(req_file)
				
				if missing_files:
					logger.error(f"压缩包缺少必要的OpenSSH文件: {', '.join(missing_files)}")
					exit(2)
			
			# 确定解压目标路径
			if has_openssh:
				target_dir = "/usr/local"
			else:
				target_dir = "/usr/local/openssh"
				os.makedirs(target_dir, exist_ok=True)
			
			# 使用zipfile解压
			try:
				with zipfile.ZipFile(self.filename) as zf:
					zf.extractall(path=target_dir)
				logger.info("解压成功")
			except Exception as e:
				logger.error(f"解压失败: {e}")
				exit(2)
		except Exception as e:
			logger.error(f"解压失败: {e}")
			exit(2)

	def extract_tar_archive(self):
		"""
		解压tar.gz格式的归档文件
		:return:
		"""
		# 检查压缩包结构
		has_openssh = False
		required_files = self.required_files
		try:
			with tarfile.open(self.filename, "r:gz") as tar:
				members = tar.getmembers()
				# 检查是否包含openssh目录
				for member in members:
					if member.name.startswith("openssh/") or member.name.startswith("./openssh/"):
						has_openssh = True
						break
				
				# 检查是否包含必要的文件
				missing_files = []
				for req_file in required_files:
					found = False
					for member in members:
						if (has_openssh and (member.name == f"openssh/{req_file}" or member.name == f"./openssh/{req_file}")) or \
						   (not has_openssh and (member.name == req_file or member.name == f"./{req_file}")):
							found = True
							break
					if not found:
						missing_files.append(req_file)
				
				if missing_files:
					logger.error(f"压缩包缺少必要的OpenSSH文件: {', '.join(missing_files)}")
					exit(2)
			
			# 确定解压目标路径
			if has_openssh:
				target_dir = "/usr/local"
			else:
				target_dir = "/usr/local/openssh"
				os.makedirs(target_dir, exist_ok=True)
			
			# 使用tarfile解压并规范化路径
			try:
				with tarfile.open(self.filename, "r:gz") as tar:
					for member in tar.getmembers():
						# 规范化路径，去除"./"前缀
						if member.name.startswith("./"):
							member.name = member.name[2:]
						# 确保路径以openssh/开头
						if has_openssh and not member.name.startswith("openssh/"):
							member.name = os.path.join("openssh", member.name)
						tar.extract(member, path=target_dir)
				logger.info("解压成功")
			except Exception as e:
				logger.error(f"解压失败: {e}")
				exit(2)
		except Exception as e:
			logger.error(f"解压失败: {e}")
			exit(2)

	def backup_directory(self, src_dir):
		"""
		备份目录到备份位置
		:param src_dir: 要备份的源目录
		:return: 备份是否成功
		"""
		if not os.path.exists(src_dir):
			logger.warning(f"源目录不存在，无需备份: [ {src_dir} ]")
			return True
			
		# 创建带时间戳的备份目录
		self.backup_dir = f"/tmp/openssh_backup_{datetime.now().strftime('%Y%m%d%H%M%S')}"
		os.makedirs(self.backup_dir, exist_ok=True)
		
		backup_cmd = f"mv {src_dir} {self.backup_dir}/"
		if os.system(backup_cmd) != 0:
			logger.error(f"备份目录失败: [ {src_dir} -> {self.backup_dir} ]")
			return False
		logger.info(f"成功备份目录: [ {src_dir} -> {self.backup_dir} ]")
		return True

	def restore(self):
		"""
		从备份恢复目录
		:return: 恢复是否成功
		"""
		if not self.backup_dir or not os.path.exists(f"{self.backup_dir}/{os.path.basename(self.ssh_install_dst)}"):
			logger.error(f"备份目录不存在，无法恢复: [ {self.backup_dir or '未设置'} ]")
			return False
			
		restore_cmd = f"rm -rf {self.ssh_install_dst} && mv {self.backup_dir}/{os.path.basename(self.ssh_install_dst)} {self.ssh_install_dst}"
		if os.system(restore_cmd) != 0:
			logger.error(f"恢复备份失败: [ {self.backup_dir} -> {self.ssh_install_dst} ]")
			return False
		logger.info(f"成功恢复备份: [ {self.backup_dir} -> {self.ssh_install_dst} ]")
		return True

	def validate_archive_format(self, filepath):
		"""
		验证压缩文件的实际格式
		:param filepath: 文件路径
		:return: 格式类型 ('zip', 'gzip', None)
		"""
		try:
			with open(filepath, 'rb') as f:
				header = f.read(4)
				if header.startswith(b'PK\x03\x04'):
					return 'zip'
				elif header.startswith(b'\x1f\x8b'):
					return 'gzip'
		except Exception as e:
			logger.error(f"验证文件格式失败: {e}")
		return None

	def extract_archive(self):
		"""
		解压归档文件的主入口函数，支持多种压缩格式
		采用整体目录替换方式
		:return:
		"""
		if not os.path.exists(self.filename):
			logger.error(f"文件不存在: [ {self.filename} ]")
			exit(2)

		# 验证文件实际格式
		actual_format = self.validate_archive_format(self.filename)
		ext = os.path.splitext(self.filename)[1].lower()
		
		# 检查扩展名和实际格式是否匹配
		if (ext == '.zip' and actual_format != 'zip') or \
		   (ext == '.gz' and actual_format != 'gzip'):
			logger.error(f"文件扩展名({ext})与实际格式({actual_format})不匹配")
			exit(2)

		# 备份原目录
		if not self.backup_directory(self.ssh_install_dst):
			exit(2)

		try:
			# 创建目标目录
			os.makedirs(self.ssh_install_dst, exist_ok=True)
			
			# 根据验证后的格式选择解压方式
			if actual_format == 'zip':
				logger.info("检测到zip格式文件，使用zipfile解压")
				with zipfile.ZipFile(self.filename) as zf:
					# 检查压缩包结构
					has_openssh = any(name.startswith(("openssh/", "./openssh/")) for name in zf.namelist())
					
					# 解压到临时目录
					temp_dir = f"{self.ssh_install_dst}_tmp"
					os.makedirs(temp_dir, exist_ok=True)
					zf.extractall(path=temp_dir)
					
					# 移动文件到目标目录
					if has_openssh:
						src_dir = os.path.join(temp_dir, "openssh")
					else:
						src_dir = temp_dir
						
					mv_cmd = f"mv {src_dir}/* {self.ssh_install_dst}/ && rm -rf {temp_dir}"
					if os.system(mv_cmd) != 0:
						raise Exception("移动解压文件失败")
						
			elif self.filename.endswith(".tar.gz"):
				logger.info("检测到tar.gz格式文件，使用tarfile解压")
				with tarfile.open(self.filename, "r:gz") as tar:
					# 检查压缩包结构
					has_openssh = any(member.name.startswith(("openssh/", "./openssh/")) for member in tar.getmembers())
					
					# 解压到目标目录
					for member in tar.getmembers():
						# 规范化路径
						if member.name.startswith("./"):
							member.name = member.name[2:]
						if has_openssh and not member.name.startswith("openssh/"):
							member.name = os.path.join("openssh", member.name)
						tar.extract(member, path=self.ssh_install_dst)
			else:
				logger.error(f"不支持的压缩格式: {self.filename}")
				raise Exception(f"Unsupported archive format: {self.filename}")
				
			logger.info("新版本解压成功")
			
		except Exception as e:
			logger.error(f"解压失败: {e}")
			
			# 恢复备份
			if not self.restore():
				logger.error("恢复备份失败！系统可能处于不一致状态")
				exit(2)
			exit(2)

	def copy(self):
		"""
		创建必要的符号链接
		:return:
		"""
		# 创建bin目录下的符号链接
		bin_list = os.listdir(self.ssh_dst_bin)
		for i in bin_list:
			src = os.path.join(self.ssh_dst_bin, i)
			dst = f"/bin/{i}"
			
			# 移除旧的链接或文件
			if os.path.lexists(dst):
				if os.path.islink(dst):
					os.remove(dst)
				elif os.path.isfile(dst):
					os.remove(dst)
				logger.info(f"已移除旧的链接/文件: {dst}")
			
			# 创建符号链接
			ln_cmd = f"ln -sf {src} {dst}"
			if os.system(ln_cmd) != 0:
				logger.error(f"创建符号链接失败: {dst}")
				exit(2)
			logger.info(f"创建符号链接成功: {dst}")

		# 创建sshd符号链接
		src_sshd = os.path.join(self.ssh_dst_sbin, "sshd")
		dst_sshd = "/usr/sbin/sshd"
		
		# 移除旧的链接或文件
		if os.path.lexists(dst_sshd):
			if os.path.islink(dst_sshd):
				os.remove(dst_sshd)
			elif os.path.isfile(dst_sshd):
				os.remove(dst_sshd)
			logger.info(f"已移除旧的sshd链接/文件: {dst_sshd}")
			
		ln_cmd = f"ln -sf {src_sshd} {dst_sshd} && chmod 755 {dst_sshd}"
		if os.system(ln_cmd) != 0:
			logger.error("创建sshd符号链接失败")
			exit(2)
		logger.info("创建sshd符号链接成功")

		logger.info("所有符号链接创建完成")

	def restart(self):
		"""
		重启SSH服务
		:return:
		"""
		if os.system("sshd -t") != 0:
			logger.error("SSH配置文件校验失败")
			exit(2)
		# 强制删除 /etc/init.d/sshd
		self.delete_init_d()
		# 重载SSH服务
		c = "systemctl daemon-reload"
		if os.system(c) != 0:
			logger.error(f"服务重载失败: {c}")
			exit(2)
		c = f"systemctl restart {self.service_new_name}"
		if os.system(c) != 0:
			logger.error(f"服务重启失败: [ {self.service_new_name} ]")
			if os.system(c) != 0:
				logger.error(f"服务重启失败: [ {self.service_new_name} ]")
				exit(2)
			else:
				logger.info(f"服务重启成功: [ {self.service_new_name} ]")
			exit(2)
		else:
			logger.info(f"服务重启成功: [ {self.service_new_name} ]")

	def version(self):
		"""
		查看版本信息
		:return:
		"""
		# 获取新版SSH版本
		new_version = getoutput("ssh -V")
		print(f"原版SSH版本: [ {self.old_version} ]")
		print(f"新版SSH版本: [ {new_version} ]")

	def status(self):
		c = f"systemctl status {self.service_new_name}"
		os.system(c)

	def open_root(self):
		"""
		开启root登录
		:return:
		"""
		# 检查是否存在root登录参数选项
		# 拼接源路径
		sshd_config = os.path.join(self.ssh_install_dst, "etc/sshd_config")
		if not os.path.exists(sshd_config):
			logger.error(f"配置文件不存在: [ {sshd_config} ]")
			exit(2)
		# 获取文件内容，判断是否存在注释符号
		text = getoutput(f"cat {sshd_config}")
		# 判断PermitRootLogin是否已配置为yes
		if text.find("#PermitRootLogin") != -1:
			# 如果存在则替换
			c = f"sed -i 's/#PermitRootLogin.*/PermitRootLogin yes/g' {sshd_config}"
			if os.system(c) != 0:
				logger.error(f"开启root登录失败[ {sshd_config} ]")
				exit(2)
			else:
				logger.info(f"开启root登录成功,请在完成升级配置后关闭[ {sshd_config} ]")
		else:
			try:
				with open(sshd_config, 'r+', encoding="utf-8") as fp:
					text = fp.read()
					fp.seek(0)
					fp.truncate()
					fp.write(text.replace("#PermitRootLogin", "PermitRootLogin"))
					fp.close()
			except Exception as e:
				logger.error(f"开启root登录失败[ {e} ]")
				print(sshd_config)
				exit(2)


	def show(self):
		"""
		显示信息
		:return:
		"""
		print(f"当前版本: [ {self.old_version} ]")
		print(f"新版链接: [ {self.url} ]")
		print(f"安装位置: [ {self.ssh_install_dst} ]")

	def check_unzip(self):
		"""
		检查是否安装unzip
		:return:
		"""
		if os.system("which unzip") != 0:
			logger.error("未安装unzip,请先安装")
			exit(2)

	def start(self, local):
		"""

		:param local: 是否使用本地文件
		:return:
		"""
		self.check_os()
		self.get_version()
		self.init_arg()
		self.show()
		if not local:
			self.download()
		self.unzip()
		self.copy()
		# sshd_shell_text()
		self.service(p="/usr/sbin/sshd")
		self.open_root()
		self.link_ssh_conf()
		self.restart()
		self.status()
		self.version()
		logger.info("升级完成")
		# 拼接路径
		etc = os.path.join(self.ssh_install_dst, "etc")
		print(f"当前服务配置文件位置: {etc}")


def use():
	"""
	使用免责声明
	:return:
	"""
	print("""
	免责声明:
	操作系统环境问题由使用者自行维护,此脚本只是个人编写，不代表任何平台或单位
	1. 本脚本仅供学习交流使用,请勿用于非法用途
	2. 脚本作者不对因使用本脚本而导致的任何问题或损失负责,请先自行阅读源码后再决定是否继续
	3. 可接受脚本bug反馈，但是不提供使用指导,请自行判断,不懂的地方可以通过-h参数获取帮助
	4. 如遇到代码问题可自行修改源码或者反馈，但是作者不会保障及时的修复
	5. 下班时间不会回复、处理任何此类脚本的问题
	""")
	while True:
		c = input("是否继续?(y/n): ")
		if c == "y":
			break
		elif c == "n":
			exit(0)
		else:
			print("输入错误,请重新输入")


def create_user():
	"""
	创建sshd用户
	:return:
	"""
	# 先判断是否已存在
	if subprocess.getstatusoutput("id sshd")[0] != 0:
		if os.system("useradd sshd") != 0:
			logger.error("用户创建失败")
			exit(2)


def service(p: str):
	"""
	创建Service
	:param p:
	:return:
	"""
	if not os.path.exists(p):
		logger.error(f"路径[ {p} ]不存在,请先安装OpenSSH或指向正确的路径")
		exit(2)
	os.system("mkdir -p /usr/lib/systemd/system_bak/")
	os.system("cp -rf /usr/lib/systemd/system/ssh*.service /usr/lib/systemd/system_bak/")
	os.system('rm -f /usr/lib/systemd/system/ssh*.service')
	if os.path.exists("/etc/init.d/sshd"):
		os.system('rm -f /etc/init.d/sshd')
	f = '/usr/lib/systemd/system/sshd.service'
	txt = f"""[Unit]
Description=This is a manually generated SSH service file
After=network.target
StartLimitIntervalSec=5

[Service]
Type=simple
User=root
ExecStart={p}
StandardOutput=journal
KillMode=control-group
Restart=on-failure
RestartPreventExitStatus=255
StandardError=inherit
SyslogIdentifier=sshd

[Install]
WantedBy=multi-user.target
Alias=sshd.service"""
	try:
		with open(f, 'w', encoding="utf-8") as fp:
			fp.write(txt)
	except Exception as e:
		logger.error(e)
		exit(2)

	os.system(f"systemctl daemon-reload")
	os.system(f"systemctl enable sshd.service")
	if os.system(f"systemctl restart sshd.service") != 0:
		logger.error("启动失败,正在恢复")
		c = "cp -rf /usr/lib/systemd/system_bak/* /usr/lib/systemd/system/"
		os.system(c)
		os.system(f"systemctl daemon-reload")
		if os.system(f"systemctl restart sshd.service") != 0:
			logger.error("启动失败,请手动启动")
			exit(2)

	os.system(f"systemctl status sshd.service")


def create_service():
	"""

	:return:
	"""
	if not os.path.isfile("/usr/local/openssh/sbin/sshd"):
		return
	logger.info("将仅生成Service文件,如需重新编译请先删除")
	service(p="/usr/sbin/sshd")
	exit(0)


if __name__ == '__main__':
	start_time = time.time()
	local_file = False
	default_url = "https://gitee.com/liumou_site/openssh/releases/download/10.0p1/OpenSSH-10.0p1-KylinV10_Tercel-arm64-OpenSSL3.5.0.tar.gz"
	arg = ArgumentParser(description='当前脚本版本: 1.0_20250102', prog="OpenSSHInstaller")
	h = f"设置下载地址(默认使用该参数),默认: {default_url}"
	arg.add_argument("-u", "--url", help=h, default=default_url, type=str)
	h = "指定本地文件,不从网页下载,例如：openssh_arm64.zip"
	arg.add_argument("-f", "--file", help=h, default=None, type=str)
	# 使用布尔值判断是否需要创建service
	arg.add_argument("-q", "--quit", help=f"创建service之后直接退出", action="store_true")
	# 解析参数
	args = arg.parse_args()
	# 获取URL值
	url = args.url
	# 获取是否使用本地文件
	file = args.file
	quit_arg = arg.parse_args().quit
	if quit_arg:
		create_service()
		exit(0)
	if file is not None:
		if not os.path.exists(file):
			logger.error(f"本地文件不存在: [ {file} ]")
			exit(2)
		else:
			local_file = True
			logger.info(f"使用本地文件: [ {file} ]")
	else:
		logger.info(f"使用默认下载地址: [ {url} ]")
		file = str(url).split("/")[-1]
	os.system("mkdir -p /var/empty")

	use()
	create_user()
	zo = OpenSSHInstaller(urls=url, file_name=file)
	zo.start(local=local_file)